import os
import torch
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import cv2
from dataloader_my import build_dataloader
from model_my import ResNetConvLSTM
from torchvision import transforms

MODEL_PATH = "./save_model/name1/best_model.pt"

DATA_DIR = r'/home/xian/mzs_project/Convlstm/dataset_fog/test'
OUTPUT_DIR = "./visualizations/name1"
os.makedirs(OUTPUT_DIR, exist_ok=True)

FRAME_LEN = 15
IMAGE_SIZE = (360, 640)
NUM_CLASSES = 2
CATEGORIES = ["nofoggy", "foggy"]  # 必须和 dataloader_my.py 一致
activations = None
gradients = None


def save_activation(module, input, output):
    """PyTorch hook: 保存前向传播的激活值"""
    global activations
    activations = output


def save_gradient(module, grad_in, grad_out):
    """PyTorch hook: 保存反向传播的梯度"""
    global gradients
    # grad_out[0] 是我们需要的梯度
    gradients = grad_out[0]


def get_grad_cam(model, x_tensor, target_class_idx):
    """
    为视频序列生成 Grad-CAM 热力图。

    Args:
        model: 你训练好的 ResNetConvLSTM 模型。
        x_tensor: 输入张量, 形状为 [1, T, C, H, W]。
        target_class_idx: 目标类别 (0=nofoggy, 1=foggy)。

    Returns:
        A list of [T] PIL Image (热力图)
        logits (模型的原始输出)
    """
    global activations, gradients
    activations = None
    gradients = None
    handle_act = model.feature_reduction.register_forward_hook(save_activation)
    handle_grad = model.feature_reduction.register_full_backward_hook(save_gradient)
    model.zero_grad()
    logits = model(x_tensor)

    score = logits[0][target_class_idx]
    score.backward()
    handle_act.remove()
    handle_grad.remove()

    if activations is None or gradients is None:
        raise RuntimeError("Hook 未能捕获激活值或梯度。请检查模型结构。")

    B_T, C_f, H_f, W_f = activations.shape

    weights = torch.mean(gradients, dim=(2, 3))

    heatmaps = []

    for t in range(B_T):  # B_T 实际上就是 T (因为 B=1)
        act_t = activations[t]  # [C_f, H_f, W_f]
        weights_t = weights[t]  # [C_f]

        # [C_f, H_f, W_f] * [C_f, 1, 1] -> 沿 C_f 维度求和 -> [H_f, W_f]
        cam_t = (act_t * weights_t.view(-1, 1, 1)).sum(dim=0)

        # 7. ReLU (只保留正贡献)
        cam_t = torch.relu(cam_t)

        # 8. 归一化到 [0, 1]
        if cam_t.max() > 0:
            cam_t = cam_t / cam_t.max()

        cam_t_np = cam_t.cpu().detach().numpy()
        cam_t_resized = cv2.resize(cam_t_np, (IMAGE_SIZE[1], IMAGE_SIZE[0]))  # (W, H)

        cam_t_heatmap_np = cv2.applyColorMap(np.uint8(255 * cam_t_resized), cv2.COLORMAP_JET)
        cam_t_heatmap = cv2.cvtColor(cam_t_heatmap_np, cv2.COLOR_BGR2RGB)  # 转为 RGB

        heatmaps.append(Image.fromarray(cam_t_heatmap))

    return heatmaps, logits


def tensor_to_pil(tensor):
    """将 [C, H, W] 张量转换为 PIL 图像"""
    return transforms.ToPILImage()(tensor.cpu())


def overlay_heatmap(original_img, heatmap_img, alpha=0.5):
    """将热力图半透明地叠加到原图上"""
    return Image.blend(original_img.convert('RGB'), heatmap_img.convert('RGB'), alpha)


def visualize():
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")

    print("加载数据...")
    test_loader = build_dataloader(
        root_dir=DATA_DIR,
        frame_len=FRAME_LEN,
        batch_size=1,
        num_workers=0,
        image_size=IMAGE_SIZE,
        mode='test'  # mode='test' 会设置 shuffle=False
    )

    print(f"加载模型: {MODEL_PATH}")
    model = ResNetConvLSTM(num_classes=NUM_CLASSES, pretrain_path=None).to(device)

    if not os.path.exists(MODEL_PATH):
        print(f"错误: 找不到模型权重 {MODEL_PATH}")
        return

    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
    model.eval()

    print("随机抽取一个样本...")
    try:
        x, y = next(iter(test_loader))
    except StopIteration:
        print("错误: 数据加载器为空，请检查测试集路径。")
        return

    x = x.to(device)

    true_label_idx = y.item()
    true_label_str = CATEGORIES[true_label_idx]
    print(f"运行 Grad-CAM... 目标类别: {true_label_str} (Idx: {true_label_idx})")

    try:
        heatmaps, logits = get_grad_cam(model, x, target_class_idx=true_label_idx)  # logits.argmax(1).item()
    except RuntimeError as e:
        print(f"Grad-CAM 失败: {e}")
        return

    pred_label_idx = logits.argmax(1).item()
    pred_label_str = CATEGORIES[pred_label_idx]

    status = "正确" if true_label_str == pred_label_str else "错误"
    print(f"样本分析 [ {status} ]: 真实标签 = {true_label_str}, 预测标签 = {pred_label_str}")

    original_frames = [tensor_to_pil(x[0, t]) for t in range(FRAME_LEN)]

    print(f"保存 {FRAME_LEN} 帧叠加图像到 {OUTPUT_DIR}...")

    grid_w, grid_h = 5, 3  # 5x3 网格
    fig, axes = plt.subplots(grid_h, grid_w, figsize=(20, 12))  # (宽, 高)
    fig.suptitle(f"Grad-CAM (True: {true_label_str}, Pred: {pred_label_str})", fontsize=16)

    axes = axes.flatten()  # 将 3x5 数组展平

    for t in range(FRAME_LEN):
        overlayed_img = overlay_heatmap(original_frames[t], heatmaps[t], alpha=0.6)
        axes[t].imshow(overlayed_img)
        axes[t].set_title(f"Frame {t}")
        axes[t].axis('off')  # 关闭坐标轴

    plt.tight_layout(rect=[0, 0.03, 1, 0.95])  # 调整布局
    plot_path = os.path.join(OUTPUT_DIR, f"grad_cam_grid (True_{true_label_str}, Pred_{pred_label_str}).png")
    plt.savefig(plot_path)

    print(f"\n可视化网格图已保存到: {plot_path}")
    print("完成。")


if __name__ == "__main__":
    visualize()